package kvraft

import (
	"bytes"
	"encoding/gob"
	"log"
	"sync"
	"sync/atomic"
	"time"

	"6.824/labgob"
	"6.824/labrpc"
	"6.824/raft"
)

const Debug = false

func DPrintf(format string, a ...interface{}) (n int, err error) {
	if Debug {
		log.Printf(format, a...)
	}
	return
}

type Op struct {
	// Your definitions here.
	// Field names must start with capital letters,
	// otherwise RPC will break.
}

type KVServer struct {
	mu      sync.Mutex
	me      int
	rf      *raft.Raft
	applyCh chan raft.ApplyMsg
	dead    int32 // set by Kill()

	maxraftstate int // snapshot if log grows this big

	// Your definitions here.
	data       map[string]string //persist
	applyIndex int               //persist
	fromIndex  int

	identifySeqs    map[int64]int64 //persist
	replyChannels   []chan int
	serverPersistCh chan bool
}

func (kv *KVServer) adjustIndex(index int) int {
	return index - kv.fromIndex
}

func (kv *KVServer) Get(args *GetArgs, reply *GetReply) {
	// Your code here.

	kv.mu.Lock()
	send_args := PutAppendArgs{args.Key, "", "Get", args.ClientId, nrand()}
	index, term, isLeader := kv.rf.Start(send_args)
	if !isLeader {
		kv.mu.Unlock()
		reply.Err = Err("wrong guy")
		return
	}

	if len(kv.replyChannels) > kv.adjustIndex(index) {
		kv.replyChannels[kv.adjustIndex(index)] = make(chan int, 1)
	} else {
		for len(kv.replyChannels)-1 < kv.adjustIndex(index-1) {
			kv.replyChannels = append(kv.replyChannels, nil)
		}
		kv.replyChannels = append(kv.replyChannels, make(chan int, 1))
	}

	toReadChannel := kv.replyChannels[kv.adjustIndex(index)]
	kv.mu.Unlock()

	select {
	case readTerm := <-toReadChannel:
		if kv.killed() {
			reply.Err = Err("killed")
			return
		}

		if readTerm != term {
			reply.Err = Err("append failed")
			return
		}

		// append success
		kv.mu.Lock()
		reply.Value = kv.data[args.Key]
		kv.mu.Unlock()
		reply.Err = ""
		return

	case <-time.After(2000 * time.Millisecond):
		if kv.killed() {
			reply.Err = Err("killed")
			return
		}

		// timeout: 2000ms
		reply.Err = Err("timeout")
		return
	}
}

func (kv *KVServer) PutAppend(args *PutAppendArgs, reply *PutAppendReply) {
	// Your code here.
	kv.mu.Lock()
	index, term, isLeader := kv.rf.Start(*args)
	if !isLeader {
		kv.mu.Unlock()
		reply.Err = Err("wrong guy")
		return
	}

	if len(kv.replyChannels) > kv.adjustIndex(index) {
		kv.replyChannels[kv.adjustIndex(index)] = make(chan int, 1)
	} else {
		for len(kv.replyChannels)-1 < kv.adjustIndex(index-1) {
			kv.replyChannels = append(kv.replyChannels, nil)
		}
		kv.replyChannels = append(kv.replyChannels, make(chan int, 1))
	}

	toReadChannel := kv.replyChannels[kv.adjustIndex(index)]
	kv.mu.Unlock()

	select {
	case applyTerm := <-toReadChannel:
		if kv.killed() {
			reply.Err = Err("killed")
			return
		}

		if applyTerm != term {
			reply.Err = Err("append failed")
			return
		}

		// append success
		reply.Err = Err("")
		return

	case <-time.After(2000 * time.Millisecond):
		if kv.killed() {
			reply.Err = Err("killed")
			return
		}

		// timeout: 2000ms
		reply.Err = Err("timeout")
		return
	}

}

func (kv *KVServer) selectApplyLogs(cmd PutAppendArgs, term int) {

	if kv.identifySeqs[cmd.ClinetId] == cmd.IdentifySeq {
		kv.applyIndex++
		if len(kv.replyChannels) > kv.adjustIndex(kv.applyIndex) && kv.replyChannels[kv.adjustIndex(kv.applyIndex)] != nil {
			kv.replyChannels[kv.adjustIndex(kv.applyIndex)] <- term
		}
		return
	}

	switch cmd.Op {
	case "Put":
		kv.data[cmd.Key] = cmd.Value
	case "Append":
		kv.data[cmd.Key] = kv.data[cmd.Key] + cmd.Value
	case "Get":
		// nothing to do
	default:
		panic("Unknown OP")
	}

	kv.identifySeqs[cmd.ClinetId] = cmd.IdentifySeq
	kv.applyIndex++
	if len(kv.replyChannels) > kv.adjustIndex(kv.applyIndex) && kv.replyChannels[kv.adjustIndex(kv.applyIndex)] != nil {
		kv.replyChannels[kv.adjustIndex(kv.applyIndex)] <- term
	}
}

func (kv *KVServer) applyLogDaemon() {
	apply_count := 0
	for msg := range kv.applyCh {
		if msg.SnapshotValid {
			kv.mu.Lock()
			kv.applySnapshotData(msg.Snapshot)
			kv.mu.Unlock()
			continue
		}

		raft.Assert(msg.CommandIndex == kv.applyIndex+1, "not msg.CommandIndex==kv.applyIndex+1")
		apply_count++

		cmd := msg.Command.(PutAppendArgs)
		term := msg.CommandTerm

		kv.mu.Lock()
		kv.selectApplyLogs(cmd, term)
		kv.mu.Unlock()

		if kv.maxraftstate != -1 && apply_count > kv.maxraftstate/100 {
			apply_count = 0
			kv.serverPersistCh <- true
		}

		if kv.killed() {
			return
		}
	}
}

func (kv *KVServer) applySnapshotData(snapshotData []byte) {
	if snapshotData == nil || len(snapshotData) < 1 {
		return
	}

	r := bytes.NewBuffer(snapshotData)
	d := labgob.NewDecoder(r)

	var data map[string]string
	var applyIndex int
	var identifySeqs map[int64]int64
	d.Decode(&data)
	d.Decode(&applyIndex)
	d.Decode(&identifySeqs)

	kv.data = data
	kv.applyIndex = applyIndex
	kv.identifySeqs = identifySeqs
	kv.fromIndex = applyIndex

	kv.replyChannels = []chan int{nil}

}

func (kv *KVServer) makeSnapshotData() []byte {
	w := new(bytes.Buffer)
	e := labgob.NewEncoder(w)

	e.Encode(kv.data)
	e.Encode(kv.applyIndex)
	e.Encode(kv.identifySeqs)

	data := w.Bytes()
	return data
}

func (kv *KVServer) doSnapshotDaemon() {
	if kv.maxraftstate == -1 {
		return
	}

	for !kv.killed() {
		select {
		case <-kv.serverPersistCh:

			kv.mu.Lock()
			if kv.rf.ZeroLogIndex() >= kv.applyIndex {
				kv.mu.Unlock()
				time.Sleep(50 * time.Millisecond)
				continue
			}

			kv.rf.Snapshot(kv.applyIndex, kv.makeSnapshotData())

			if kv.adjustIndex(kv.applyIndex) < len(kv.replyChannels) {
				kv.replyChannels = kv.replyChannels[kv.adjustIndex(kv.applyIndex):]
			} else {
				kv.replyChannels = []chan int{nil}
			}

			kv.fromIndex = kv.applyIndex
			kv.mu.Unlock()
			time.Sleep(50 * time.Millisecond)

		}
	}
}

// the tester calls Kill() when a KVServer instance won't
// be needed again. for your convenience, we supply
// code to set rf.dead (without needing a lock),
// and a killed() method to test rf.dead in
// long-running loops. you can also add your own
// code to Kill(). you're not required to do anything
// about this, but it may be convenient (for example)
// to suppress debug output from a Kill()ed instance.
func (kv *KVServer) Kill() {
	atomic.StoreInt32(&kv.dead, 1)
	kv.rf.Kill()
	// Your code here, if desired.
}

func (kv *KVServer) killed() bool {
	z := atomic.LoadInt32(&kv.dead)
	return z == 1
}

// servers[] contains the ports of the set of
// servers that will cooperate via Raft to
// form the fault-tolerant key/value service.
// me is the index of the current server in servers[].
// the k/v server should store snapshots through the underlying Raft
// implementation, which should call persister.SaveStateAndSnapshot() to
// atomically save the Raft state along with the snapshot.
// the k/v server should snapshot when Raft's saved state exceeds maxraftstate bytes,
// in order to allow Raft to garbage-collect its log. if maxraftstate is -1,
// you don't need to snapshot.
// StartKVServer() must return quickly, so it should start goroutines
// for any long-running work.
func StartKVServer(servers []*labrpc.ClientEnd, me int, persister *raft.Persister, maxraftstate int) *KVServer {
	// call labgob.Register on structures you want
	// Go's RPC library to marshall/unmarshall.
	labgob.Register(Op{})

	kv := new(KVServer)
	kv.me = me
	kv.maxraftstate = maxraftstate

	// You may need initialization code here.

	kv.applyCh = make(chan raft.ApplyMsg)
	kv.rf = raft.Make(servers, me, persister, kv.applyCh)

	// You may need initialization code here.
	gob.Register(PutAppendArgs{})
	kv.data = make(map[string]string, 0)
	kv.applyIndex = kv.rf.ZeroLogIndex()
	kv.fromIndex = kv.rf.ZeroLogIndex()
	kv.replyChannels = append(kv.replyChannels, nil)
	kv.identifySeqs = make(map[int64]int64)
	kv.applySnapshotData(kv.rf.SnapshotForServer())

	if kv.maxraftstate != -1 {
		kv.serverPersistCh = make(chan bool, 1)
	}

	go kv.applyLogDaemon()
	go kv.doSnapshotDaemon()

	return kv
}
